import os
import numpy as np
import paddle
import paddle.nn as nn
import paddle.optimizer as optim
from paddle.io import Dataset, DataLoader
from paddle.vision import transforms
from paddle.regularizer import L2Decay
import cv2
from PIL import Image
import random

# 数据集类 - 处理工业电池图像数据
class BatteryDataset(Dataset):
    def __init__(self, data_dir, mode='train', transform=None):
        self.data_dir = data_dir
        self.mode = mode
        self.transform = transform
        self.classes = os.listdir(data_dir)
        self.class_to_idx = {cls: i for i, cls in enumerate(self.classes)}
        self.samples = self._load_samples()
        
    def _load_samples(self):
        samples = []
        for cls in self.classes:
            cls_dir = os.path.join(self.data_dir, cls)
            if not os.path.isdir(cls_dir):
                continue
            for img_name in os.listdir(cls_dir):
                img_path = os.path.join(cls_dir, img_name)
                if os.path.isfile(img_path) and img_path.endswith(('.jpg', '.png', '.jpeg')):
                    samples.append((img_path, self.class_to_idx[cls]))
        return samples
    
    def __getitem__(self, idx):
        img_path, label = self.samples[idx]
        image = Image.open(img_path).convert('RGB')
        
        if self.transform:
            image = self.transform(image)
        
        return image, label
    
    def __len__(self):
        return len(self.samples)

# 质检模型 - 基于ResNet架构
class BatteryInspectionModel(nn.Layer):
    def __init__(self, num_classes=2, pretrained=True):
        super(BatteryInspectionModel, self).__init__()
        # 使用飞桨预训练的ResNet50作为特征提取器
        self.backbone = paddle.vision.models.resnet50(pretrained=pretrained)
        
        # 修改最后的全连接层以适应我们的分类任务
        self.backbone.fc = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(self.backbone.fc.in_features, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )
    
    def forward(self, x):
        return self.backbone(x)

# 数据增强和预处理
def get_transforms(mode='train'):
    if mode == 'train':
        return transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.RandomHorizontalFlip(0.5),
            transforms.RandomVerticalFlip(0.5),
            transforms.RandomRotation(15),
            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
    else:
        return transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

# 模型训练函数
def train_model(model, train_loader, val_loader, num_epochs=10, lr=0.001):
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(
        parameters=model.parameters(),
        learning_rate=lr,
        weight_decay=L2Decay(1e-4)
    )
    
    # 学习率调度器
    scheduler = optim.lr.StepDecay(lr, step_size=5, gamma=0.1)
    
    # 训练循环
    for epoch in range(num_epochs):
        # 训练阶段
        model.train()
        train_loss = 0.0
        correct = 0
        total = 0
        
        for batch_idx, (inputs, targets) in enumerate(train_loader):
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            
            optimizer.clear_grad()
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.equal(targets).sum().item()
            
            if (batch_idx + 1) % 10 == 0:
                print(f'Epoch: {epoch+1}/{num_epochs} | Batch: {batch_idx+1}/{len(train_loader)} | '
                      f'Loss: {train_loss/(batch_idx+1):.4f} | Acc: {100.*correct/total:.2f}%')
        
        # 验证阶段
        model.eval()
        val_loss = 0.0
        val_correct = 0
        val_total = 0
        
        with paddle.no_grad():
            for inputs, targets in val_loader:
                outputs = model(inputs)
                loss = criterion(outputs, targets)
                
                val_loss += loss.item()
                _, predicted = outputs.max(1)
                val_total += targets.size(0)
                val_correct += predicted.equal(targets).sum().item()
        
        print(f'Epoch: {epoch+1}/{num_epochs} | Val Loss: {val_loss/len(val_loader):.4f} | '
              f'Val Acc: {100.*val_correct/val_total:.2f}%')
        
        # 更新学习率
        scheduler.step()
    
    return model

# 模型评估函数
def evaluate_model(model, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    total = 0
    criterion = nn.CrossEntropyLoss()
    
    with paddle.no_grad():
        for inputs, targets in test_loader:
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            
            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.equal(targets).sum().item()
    
    print(f'Test Loss: {test_loss/len(test_loader):.4f} | Test Acc: {100.*correct/total:.2f}%')
    
    # 计算精确率、召回率和F1分数
    from sklearn.metrics import classification_report
    all_targets = []
    all_predicted = []
    
    with paddle.no_grad():
        for inputs, targets in test_loader:
            outputs = model(inputs)
            _, predicted = outputs.max(1)
            all_targets.extend(targets.numpy())
            all_predicted.extend(predicted.numpy())
    
    print(classification_report(all_targets, all_predicted, target_names=['合格', '不合格']))

# 模型推理函数 - 单张图像预测
def predict_single_image(model, image_path, transform):
    model.eval()
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0)  # 添加batch维度
    
    with paddle.no_grad():
        output = model(image)
        probs = paddle.nn.functional.softmax(output, axis=1)
        confidence, pred_class = paddle.max(probs, axis=1)
    
    return pred_class.item(), confidence.item()

# 主函数
def main():
    # 数据路径
    data_dir = 'path/to/your/dataset'  # 替换为实际数据路径
    
    # 创建数据集
    train_dataset = BatteryDataset(data_dir, mode='train', transform=get_transforms('train'))
    val_dataset = BatteryDataset(data_dir, mode='val', transform=get_transforms('val'))
    test_dataset = BatteryDataset(data_dir, mode='test', transform=get_transforms('test'))
    
    # 创建数据加载器
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)
    
    # 创建模型
    model = BatteryInspectionModel(num_classes=2)
    
    # 训练模型
    model = train_model(model, train_loader, val_loader, num_epochs=10)
    
    # 评估模型
    evaluate_model(model, test_loader)
    
    # 保存模型
    paddle.save(model.state_dict(), 'battery_inspection_model.pdparams')
    
    # 示例：单张图像预测
    sample_image = 'path/to/sample/image.jpg'  # 替换为实际图像路径
    pred_class, confidence = predict_single_image(model, sample_image, get_transforms('test'))
    print(f'预测结果: {"合格" if pred_class == 0 else "不合格"}, 置信度: {confidence:.4f}')

if __name__ == '__main__':
    main()    